from PreRec.BPRMF.Recommender import Recommender as BPRRecommender
from PreRec.POP.Recommender import Recommender as POPRecommender
from PreRec.ICF.Recommender import Recommender as ICFRecommender
from PreRec.UCF.Recommender import Recommender as UCFRecommender
from typing import Dict, List
from torch.cuda import is_available
from tqdm import tqdm
import time
import torch
from utils.log_output import simple_text_log


class Recommender:
    def __init__(self, hyper_params):
        self.hyper_params = hyper_params

        simple_text_log('train', 'BPR Training...')
        if 'bpr' in hyper_params['methods']:
            self.bpr_recommender = BPRRecommender(hyper_params)
        if 'pop' in hyper_params['methods']:
            self.pop_recommender = POPRecommender(hyper_params)
        simple_text_log('train', 'ICF Initializing...')
        if 'icf' in hyper_params['methods']:
            self.icf_recommender = ICFRecommender(hyper_params)
        simple_text_log('train', 'UCF Initializing...')
        if 'ucf' in hyper_params['methods']:
            self.ucf_recommender = UCFRecommender(hyper_params)

        # POP Recommender are used for filling the remain items
        self.bpr_cnt = hyper_params['bpr_cnt']
        self.icf_cnt = hyper_params['icf_cnt']
        self.ucf_cnt = hyper_params['ucf_cnt']
        self.pop_cnt = hyper_params['pop_cnt']

    def recommend(self) -> Dict[int, List[int]]:
        simple_text_log('train', 'Generating PreRec results...')
        if 'bpr' in self.hyper_params['methods']:
            bpr_result = self.bpr_recommender.recommend(self.bpr_cnt)
        if 'icf' in self.hyper_params['methods']:
            icf_result = self.icf_recommender.recommend(self.icf_cnt)
        if 'ucf' in self.hyper_params['methods']:
            ucf_result = self.ucf_recommender.recommend(self.ucf_cnt)
        if 'pop' in self.hyper_params['methods']:
            pop_result = self.pop_recommender.recommend(self.pop_cnt)

        result: Dict[int, List[int]] = {}

        print('Recommender: Merging results...')
        for user in tqdm(range(self.hyper_params['user_cnt'])):
            cur_result = []

            if 'bpr' in self.hyper_params['methods'] and user in bpr_result:
                cnt = self.hyper_params['methods']['bpr']
                for rec in bpr_result[user]:
                    if rec in cur_result: continue
                    cur_result.append(rec)
                    cnt -= 1

                    if cnt == 0:
                        break
            if 'icf' in self.hyper_params['methods'] and user in icf_result:
                cnt = self.hyper_params['methods']['icf']
                for rec in icf_result[user]:
                    if rec in cur_result: continue
                    cur_result.append(rec)
                    cnt -= 1

                    if cnt == 0:
                        break
            if 'ucf' in self.hyper_params['methods'] and user in ucf_result:
                cnt = self.hyper_params['methods']['ucf']
                for rec in ucf_result[user]:
                    if rec in cur_result: continue
                    cur_result.append(rec)
                    cnt -= 1

                    if cnt == 0:
                        break
            if 'pop' in self.hyper_params['methods']:
                cnt = 500 - len(cur_result)
                for rec in pop_result:
                    if rec in cur_result: continue
                    cur_result.append(rec)
                    cnt -= 1

                    if cnt == 0:
                        break

            result[user] = cur_result

        return result


if __name__ == '__main__':
    hyper_params = {
        'datasets': 'datasets/data_small',
        'device': 'cuda' if is_available() else 'cpu',
        'batch_size': 128 if is_available() else 64,
        'epochs': 5 if is_available() else 2,
        'embed_dim': 64,
        'neg_sample_cnt': 300,
        'epsilon': 0.1,
        'bpr_cnt': 500,
        'icf_cnt': 500,
        'ucf_cnt': 500,
        'pop_cnt': 500,
        'methods': {'icf': 60, 'bpr': 60, 'pop': 85, 'ucf': 295}
    }

    print('Recommender: Initializing...')
    recommender = Recommender(hyper_params)
    print('Recommender: Finished.')

    start_time = time.time()
    print('Recommender: Start recommend...')
    result = recommender.recommend()
    print(f'Recommender: Finished. Time: {time.time()-start_time} seconds.')
    print(len(result[0]))

    # Save result
    result_tensor = torch.empty([len(result), 500], dtype=torch.long)
    for i in range(len(result)):
        result_tensor[i] = torch.tensor(result[i])

    torch.save(result_tensor, 'prerec_result_small.pkl')
